#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
正式训练脚本
使用优化后的配置进行完整的Potsdam模型训练
"""

import os
import sys
import json
import torch
import random
import argparse
import numpy as np
import time
import warnings
from pathlib import Path

# 抑制警告信息
warnings.filterwarnings("ignore")
os.environ['PYTHONWARNINGS'] = 'ignore'

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from lib.utils.tools.logger import Logger as Log
from lib.utils.tools.configer import Configer


def main():
    """主函数"""
    print("🚀 Potsdam正式训练开始")
    print("=" * 80)
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Potsdam正式训练')
    parser.add_argument('--config', default='configs/potsdam/H_48_D_4_proto_test.json',
                       help='配置文件路径')
    parser.add_argument('--max-iters', type=int, default=100000,
                       help='最大训练迭代数')
    parser.add_argument('--display-interval', type=int, default=50,
                       help='显示间隔')
    parser.add_argument('--test-interval', type=int, default=2000,
                       help='验证间隔')
    parser.add_argument('--save-interval', type=int, default=2000,
                       help='保存间隔')
    args = parser.parse_args()
    
    # 设置随机种子
    seed = 304
    random.seed(seed)
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    # 创建配置器
    configer = Configer(configs=args.config)
    
    # 自动检测并设置最佳设备
    if torch.backends.mps.is_available():
        configer.update(['gpu'], [0])  # 使用MPS
        print("🚀 检测到MPS，启用GPU加速")
    elif torch.cuda.is_available():
        configer.update(['gpu'], [0])  # 使用CUDA
        print("🚀 检测到CUDA，启用GPU加速")
    else:
        configer.update(['gpu'], [])  # 使用CPU
        print("⚠️  未检测到GPU，使用CPU模式")
    
    # 更新训练参数
    configer.update(['solver', 'max_iters'], args.max_iters)
    configer.update(['solver', 'display_iter'], args.display_interval)
    configer.update(['solver', 'test_interval'], args.test_interval)
    configer.update(['checkpoints', 'save_iters'], args.save_interval)
    
    # 设置日志级别
    Log.init(logfile_level="info",
             stdout_level="info",
             log_file=configer.get('logging', 'log_file'),
             log_format="%(asctime)s %(levelname)-7s %(message)s",
             rewrite=True)
    
    # 打印训练配置信息
    print(f"📊 数据集: {configer.get('dataset')}")
    print(f"🏗️  模型: {configer.get('network', 'model_name')}")
    print(f"📚 学习率: {configer.get('lr', 'base_lr')}")
    print(f"🔄 最大迭代: {args.max_iters}")
    print(f"📈 显示间隔: {args.display_interval}")
    print(f"🔍 验证间隔: {args.test_interval}")
    print(f"💾 保存间隔: {args.save_interval}")
    print(f"👥 工作进程: {configer.get('data', 'workers')}")
    print(f"📦 训练批次大小: {configer.get('train', 'batch_size')}")
    print(f"📦 验证批次大小: {configer.get('val', 'batch_size')}")
    print("=" * 80)
    
    # 创建检查点目录
    checkpoint_dir = configer.get('checkpoints', 'checkpoints_dir')
    os.makedirs(checkpoint_dir, exist_ok=True)
    print(f"📁 检查点目录: {checkpoint_dir}")
    
    # 估算训练时间
    estimated_time = args.max_iters * 1.5 / 60  # 假设每次迭代1.5秒（考虑批次增大）
    print(f"⏱️  预计训练时间: {estimated_time:.1f}分钟 ({estimated_time/60:.1f}小时)")
    print("=" * 80)
    
    try:
        # 创建简化版训练器
        from simplified_trainer import SimplifiedTrainer
        trainer = SimplifiedTrainer(configer)
        
        print("✅ 训练器创建成功，开始正式训练...")
        print("🎯 训练目标: 完成Potsdam语义分割模型训练")
        print()
        
        # 记录开始时间
        start_time = time.time()
        
        # 开始训练
        trainer.train()
        
        # 计算总训练时间
        total_time = time.time() - start_time
        
        print("\n" + "=" * 80)
        print("🎉 正式训练完成!")
        print(f"⏱️  总训练时间: {total_time/60:.1f}分钟 ({total_time/3600:.1f}小时)")
        print(f"📊 最佳验证损失: {trainer.best_val_loss:.4f}")
        print(f"💾 模型保存在: {checkpoint_dir}")
        print("=" * 80)
        
        # 生成训练报告
        generate_training_report(trainer, total_time, args)
        
    except Exception as e:
        print(f"\n❌ 训练失败: {str(e)}")
        import traceback
        traceback.print_exc()


def generate_training_report(trainer, total_time, args):
    """生成训练报告"""
    report_content = f"""# Potsdam模型训练报告

## 训练配置
- 数据集: Potsdam
- 模型: HRNet-W48 + ProtoSeg
- 最大迭代数: {args.max_iters}
- 训练批次大小: {trainer.configer.get('train', 'batch_size')}
- 验证批次大小: {trainer.configer.get('val', 'batch_size')}
- 学习率: {trainer.configer.get('lr', 'base_lr')}

## 训练结果
- 总训练时间: {total_time/60:.1f}分钟 ({total_time/3600:.1f}小时)
- 最佳验证损失: {trainer.best_val_loss:.4f}
- 早停计数: {trainer.patience_counter}/{trainer.patience}

## 性能统计
- 平均每次迭代时间: {total_time/args.max_iters:.3f}秒
- 训练效率: {args.max_iters/(total_time/60):.1f} iterations/min

## 文件位置
- 检查点目录: {trainer.configer.get('checkpoints', 'checkpoints_dir')}
- 日志文件: {trainer.configer.get('logging', 'log_file')}

## 训练时间
开始时间: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time() - total_time))}
结束时间: {time.strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    report_path = "training_report.md"
    with open(report_path, 'w', encoding='utf-8') as f:
        f.write(report_content)
    
    print(f"📝 训练报告已生成: {report_path}")


if __name__ == "__main__":
    main()
